/*********************************************************************************
 *      Copyright:  (C) 2024 ioter<ioter@gmail.com>
 *                  All rights reserved.
 *
 *       Filename:  proxy.c
 *    Description:  This file 
 *                 
 *        Version:  1.0.0(30/06/24)
 *         Author:  ioter <ioter@gmail.com>
 *      ChangeLog:  1, Release initial version on "30/06/24 19:17:12"
 *                 
 ********************************************************************************/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <openssl/ssl.h>
#include <openssl/err.h>
#include <pthread.h>
#include <sys/resource.h>
#include <sys/time.h>
#include <fcntl.h>

#define PROXY_PORT   8899
#define SERVER_PORT  4433
#define SERVER_HOST  "127.0.0.1"

typedef struct 
{
    int client_sock;
    SSL_CTX *ctx;
    struct sockaddr_in client_addr;
} client_info_t;

void proxy_client(int client_sock, SSL_CTX *ctx, struct sockaddr_in *client_addr);
void *client_handler(void *arg);
void set_max_fd_limit();

void initialize_openssl()
{
	SSL_load_error_strings();
	OpenSSL_add_ssl_algorithms();
}

void cleanup_openssl()
{
	EVP_cleanup();
}

SSL_CTX *create_context()
{
	const SSL_METHOD *method;
	SSL_CTX *ctx;

	method = SSLv23_client_method();

	ctx = SSL_CTX_new(method);
	if(!ctx)
	{
		perror("Uable to create SSL	context");
		ERR_print_errors_fp(stderr);
		exit(EXIT_FAILURE);
	}

	return ctx;
}


void set_max_fd_limit()
{
	struct rlimit limit;
	limit.rlim_cur = 10240;
	limit.rlim_max = 10240;
	if(setrlimit(RLIMIT_NOFILE, &limit) == -1)
	{
		perror("setrlimit");
		exit(EXIT_FAILURE);
	}
}


void proxy_client(int client_sock, SSL_CTX *ctx, struct sockaddr_in *client_addr)
{
	int					server_sock;
	struct sockaddr_in server_addr;

	if((server_sock= socket(AF_INET, SOCK_STREAM, 0)) < 0)
	{
		perror("socket create failed");
		exit(EXIT_FAILURE);
		return;
	}

	struct timeval timeout;
    timeout.tv_sec = 10; 
    timeout.tv_usec = 0;

    if (setsockopt(server_sock, SOL_SOCKET, SO_RCVTIMEO, (char *)&timeout, sizeof(timeout)) < 0) 
	{
        perror("setsockopt failed");
    }

    if (setsockopt(server_sock, SOL_SOCKET, SO_SNDTIMEO, (char *)&timeout, sizeof(timeout)) < 0) 
	{
        perror("setsockopt failed");
    }

	server_addr.sin_family = AF_INET;
	server_addr.sin_port = htons(SERVER_PORT);
	if(inet_pton(AF_INET, SERVER_HOST, &server_addr.sin_addr) <= 0)
	{
		perror("Address not supported");
		close(server_sock);
		close(client_sock);
		return;
	}
	printf("Address converted\n");

	if(connect(server_sock, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0)
	{
		perror("connect to server failed");
		close(server_sock);
		close(client_sock);
		return;
	}
	printf("Connected to server\n");

	if (!ctx) 
	{
        printf("SSL context is NULL\n");
    } 
	else 
	{
        printf("SSL context is valid\n");
    }

	SSL *ssl = SSL_new(ctx);
	if(!ssl)
    {
        perror("Unable to create SSL structure");
        ERR_print_errors_fp(stderr);
		close(server_sock);
		close(client_sock);
        return;
    }


	if(SSL_set_fd(ssl, server_sock) == 0)
	{
		printf("set failede\n");

		ERR_print_errors_fp(stderr);
        SSL_free(ssl);
        close(server_sock);
		close(client_sock);
        return;
	}

	initialize_openssl();
	int ssl_connect_result;
	while((ssl_connect_result = SSL_connect(ssl)) != 1)
	{
		int ssl_error = SSL_get_error(ssl, ssl_connect_result);
		if(ssl_error == SSL_ERROR_WANT_READ || ssl_error == SSL_ERROR_WANT_WRITE)
		{
			fd_set fds;
			struct timeval tv;
			FD_ZERO(&fds);
			FD_SET(server_sock, &fds);
			tv.tv_sec = 5;
			tv.tv_usec = 0;

			int ret =select (server_sock + 1, ssl_error == SSL_ERROR_WANT_READ ? &fds: NULL, 
								SSL_ERROR_WANT_WRITE ? &fds: NULL, NULL, &tv);

				if(ret <= 0)
				{
					perror("select failed");
					SSL_free(ssl);
					close(server_sock);
					close(client_sock);
					return;
				}
				
		}
		else
		{
			printf("SSL_connect failed with error code %d\n", ssl_error);
            ERR_print_errors_fp(stderr);
	        SSL_free(ssl);
   	 	    close(server_sock);
        	close(client_sock);
        	return;

		}
	}


	char buffer[1024];
	printf("Entering while loop\n");

	while(1)
	{
		printf("proxy_client: Waiting to read from client port %d\n", ntohs(client_addr->sin_port));
		int rv = read(client_sock, buffer, sizeof(buffer));
		if(rv > 0)
		{
			buffer[rv] = 0;
			printf("Read from client port %d: %s\n", ntohs(client_addr->sin_port), buffer);

        	if(SSL_write(ssl, buffer, rv) <= 0)
			{
				ERR_print_errors_fp(stderr);
				break;
			}
			else
			{
				printf("Message sent to server\n");
			}
			int bytes = SSL_read(ssl, buffer, sizeof(buffer));
			if(bytes > 0)
			{
				buffer[bytes] = 0;
				printf("Read from server: %s\n", buffer);
				write(client_sock, buffer, bytes);
				printf("Message sent back to client\n");
			}
			else if(bytes < 0)
			{
				ERR_print_errors_fp(stderr);	
				break;
			}
			else
			{
				printf("No message read from server\n");
				break;
			}
		}
		else if(rv < 0)
		{
			if(errno == EAGAIN || errno == EWOULDBLOCK)
			{
				continue;
			}

			perror("Read from client failed");
			break;
		}
		else
		{
			printf("No data read from client\n");
			break;
		}
	}

	SSL_shutdown(ssl);
	SSL_free(ssl);
	close(server_sock);
	close(client_sock);	
}

void *client_handler(void *arg)
{
	printf("client_handler: Started\n");
	client_info_t *client_info = (client_info_t *)arg;
	proxy_client(client_info->client_sock, client_info->ctx, &client_info->client_addr);
	close(client_info->client_sock);
	free(client_info);
	printf("client_handler: Finished\n");
	pthread_exit(NULL);
}

int main()
{
	set_max_fd_limit();
	SSL_CTX * ctx;
	initialize_openssl();
	ctx = create_context();

	int proxy_sock;		
	struct sockaddr_in proxy_addr;
	struct sockaddr_in client_addr;
	socklen_t client_addr_len = sizeof(client_addr);

	proxy_sock = socket(AF_INET, SOCK_STREAM, 0);
	if(proxy_sock < 0)
	{
		printf("create socket failed");
		exit(EXIT_FAILURE);
	}

	proxy_addr.sin_family = AF_INET;
	proxy_addr.sin_port = htons(PROXY_PORT);
	proxy_addr.sin_addr.s_addr = INADDR_ANY;

	if(bind(proxy_sock, (struct sockaddr *)&proxy_addr, sizeof(proxy_addr)) < 0)
	{
		perror("bind failed");
		close(proxy_sock);
		exit(EXIT_FAILURE);
	}

	if(listen(proxy_sock, 1024) < 0)
	{
		perror("listen failed");
		close(proxy_sock);
		exit(EXIT_FAILURE);
	}

	printf("Proxy listening on port: %d\n", PROXY_PORT);

	while(1)
	{
		int client_sock = accept(proxy_sock, (struct sockaddr *)&client_addr, &client_addr_len);
		if(client_sock < 0)
		{
			perror("accept failed");
			continue;
		}

		printf("main: New connection accepted\n");

		client_info_t *client_info = (client_info_t *)malloc(sizeof(client_info_t));
        if (!client_info) 
		{
            perror("Malloc failed");
            close(client_sock);
            continue;
        }

		client_info->client_sock = client_sock;
        client_info->ctx = ctx;
        client_info->client_addr = client_addr;

		pthread_t						tid;
		if(pthread_create(&tid, NULL, client_handler, client_info) != 0)
		{
			perror("pthread_create failed");
			close(client_sock);
            free(client_info);
            continue;
		}
		pthread_detach(tid);

	}

	close(proxy_sock);
	SSL_CTX_free(ctx);
	cleanup_openssl();

	return 0;
}
